#include <mambo.hpp>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <dns_ops.h>
#include <iostream>
using namespace std;

int main(int argc, char** argv)
{
  mpi::communicator mycomm(argc,argv);

  typedef  double real_type;
  typedef  complex<double> complex_type;
  const int n_var=3;
  int niters=2;
  int n=12288;
  int npx=512;
  int npy=512;
  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-s")<a.size()) //get the z
      n=atoi(argv[++ic]);
    else if(a.find("-npx")<a.size()) //x-partition
      npx=atoi(argv[++ic]);
    else if(a.find("-npy")<a.size()) //x-partition
      npy=atoi(argv[++ic]);
    else if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    else if(a.find("-h")<a.size()) 
      cerr << "Usage: -s fft-size -npx xgrid -npy ygrid -i iterations " << std::endl;
    ++ic;
  }

  int nx=n/npx;
  int ny=n/npy;
  int nz=n;
  int ny_half=ny/2;
  int ntot=nx*ny*nz*3;

  int np=omp_get_max_threads();
  vector<real_type*> A(np);
  vector<real_type*> B(np);
#if defined(USE_MALLOC)
  int *kx_map=static_cast<int*>(malloc(nx*sizeof(int)));
  int *ky_map=static_cast<int*>(malloc(ny_half*sizeof(int)));
  int *kz_map=static_cast<int*>(malloc(nz*sizeof(int)));
#else
  int *kx_map=new int[nx];
  int *ky_map=new int[ny_half];
  int *kz_map=new int[nz];
#endif
  for(int i=0; i<nx; ++i) kx_map[i]=i+mycomm.rank()*nx;
  for(int i=0; i<ny_half; ++i) ky_map[i]=i+mycomm.rank()*ny;
  for(int i=0; i<nz; ++i) kz_map[i]=i;

#pragma omp parallel 
  {
    int ip=omp_get_thread_num();
#if defined(USE_MALLOC)
    A[ip]=static_cast<real_type*>(malloc(ntot*sizeof(real_type)));
    B[ip]=static_cast<real_type*>(malloc(ntot*sizeof(real_type)));
#else
    A[ip]=new real_type[ntot];
    B[ip]=new real_type[ntot];
#endif
    std::fill(A[ip],A[ip]+ntot,0.0);
    std::fill(B[ip],B[ip]+ntot,0.0);
  }

  cout.setf(ios::fixed, ios::floatfield);
  cout.precision(5);
  cout << "BW nx_loc=" << nx << " ny_loc=" << ny << " nz_loc=" << nz << endl;
  cout.flush();

  mycomm.barrier();
  if(mycomm.rank()==0) cout << "PHASEC" << endl;
  //mycomm.barrier();
  Timer wtimer;
  vector<double> t(omp_get_max_threads()*4,0.0);
#pragma omp parallel 
  {
    int ip=omp_get_thread_num();
    int nodeid=mycomm.rank();
    double timing[]={0.0,0.0,0.0,0.0};
    Timer myclock;
    int it_max=niters/2;
    for(int i=0; i<it_max; ++i)
    {
#if defined(VOR_3XYZ_C)
      dns_ops::vorticity_3xyz(A[ip],B[ip],kx_map,nx,ky_map,ny_half,kz_map,nz);
      dns_ops::vorticity_3xyz(B[ip],A[ip],kx_map,nx,ky_map,ny_half,kz_map,nz);
#endif
#if defined(VOR_3XYZ_F)
      dns_ops::vorticity_3xyz_f(complex_ptr(A[ip]),complex_ptr(B[ip]),kx_map,nx,ky_map,ny_half,kz_map,nz);
      dns_ops::vorticity_3xyz_f(complex_ptr(B[ip]),complex_ptr(A[ip]),kx_map,nx,ky_map,ny_half,kz_map,nz);
#endif
#if defined(VOR_XY3Z_C)
      dns_ops::vorticity_xy3z(A[ip],B[ip],kx_map,nx,ky_map,ny_half,kz_map,nz);
      dns_ops::vorticity_xy3z(B[ip],A[ip],kx_map,nx,ky_map,ny_half,kz_map,nz);
#endif
#if defined(VOR_XY3Z_F)
      dns_ops::vorticity_xy3z_f(complex_ptr(A[ip]),complex_ptr(B[ip]),kx_map,nx,ky_map,ny_half,kz_map,nz);
      dns_ops::vorticity_xy3z_f(complex_ptr(B[ip]),complex_ptr(A[ip]),kx_map,nx,ky_map,ny_half,kz_map,nz);
#endif
#if defined(CVOR_XYZ3)
      dns_ops::vorticity_xyz3(A[ip],B[ip],kx_map,nx,ky_map,ny_half,kz_map,nz);
      dns_ops::vorticity_xyz3(B[ip],A[ip],kx_map,nx,ky_map,ny_half,kz_map,nz);
#endif
#if defined(FVOR_XYZ3)
      dns_ops::vorticity_xyz3_f(complex_ptr(A[ip]),complex_ptr(B[ip]),kx_map,nx,ky_map,ny_half,kz_map,nz);
      dns_ops::vorticity_xyz3_f(complex_ptr(B[ip]),complex_ptr(A[ip]),kx_map,nx,ky_map,ny_half,kz_map,nz);
#endif
    }
  }
  double tot_time=wtimer.elapsed();
  mycomm.barrier();
  if(mycomm.rank() == 0) cout << "PHASED" << endl;

  //for(int ip=1,p4=4; ip<np; ++ip,p4+=4)
  //{
  //  t[0]+=t[p4];
  //  t[1]+=t[p4+1];
  //  t[2]+=t[p4+2];
  //}

  double nodefac= 1/static_cast<double>(niters*omp_get_max_threads());
  //using 4 for RK4
  cout << "BW OMP_NUM_THREADS= " << omp_get_max_threads() 
    << " MPIRANK= " << mycomm.rank()
    <<  " wtime= " << tot_time 
    << "  wtime/iter= " << tot_time/static_cast<double>(niters)
    << "  wtime/step= " << 4.0*tot_time/static_cast<double>(niters)
    << endl;

#if defined(USE_MALLOC)
  for(int ip=0; ip<np; ++ip) free(A[ip]);
  for(int ip=0; ip<np; ++ip) free(B[ip]);
  free(kx_map);
  free(ky_map);
  free(kz_map);
#else
  for(int ip=0; ip<np; ++ip) delete [] A[ip];
  for(int ip=0; ip<np; ++ip) delete [] B[ip];
  delete [] kx_map;
  delete [] ky_map;
  delete [] kz_map;
#endif

  return 0;
}

